"""
模拟脑波数据生成器

本文件用于生成模拟的脑波数据，用于测试脑波监测系统的各个接口。
主要功能：
1. 生成连续30秒的模拟脑波数据
2. 模拟真实的脑波信号特征（不同频段的波形）
3. 支持多通道数据生成
4. 提供数据格式化功能，适配系统接口

原理：
- 使用正弦波叠加模拟不同频段的脑波信号
- Delta波(0.5-4Hz)、Theta波(4-8Hz)、Alpha波(8-13Hz)、Beta波(13-30Hz)、Gamma波(30-100Hz)
- 添加随机噪声模拟真实环境
- 按照系统要求的数据格式进行封装
"""

import math
import time
from datetime import datetime, timezone
from typing import List, Dict, Any
import json
import random
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from app.config import settings


class MockEEGDataGenerator:
    """模拟脑波数据生成器"""
    
    def __init__(self, sampling_rate: int = 250, channels: int = 8):
        """
        初始化数据生成器
        
        Args:
            sampling_rate: 采样率，默认250Hz
            channels: 通道数，默认8通道
        """
        self.sampling_rate = sampling_rate
        self.channels = channels
        self.channel_names = [f"CH{i+1}" for i in range(channels)]
        
        # 脑波频段定义
        self.frequency_bands = {
            'delta': (0.5, 4),      # Delta波
            'theta': (4, 8),        # Theta波
            'alpha': (8, 13),       # Alpha波
            'beta': (13, 30),       # Beta波
            'gamma': (30, 50)       # Gamma波
        }
    
    def generate_brain_wave_signal(self, duration: float, base_freq: float, amplitude: float = 1.0) -> List[float]:
        """
        生成指定频率的脑波信号
        
        Args:
            duration: 信号持续时间（秒）
            base_freq: 基础频率（Hz）
            amplitude: 振幅
            
        Returns:
            信号数据点列表
        """
        samples = int(duration * self.sampling_rate)
        signal = []
        
        for i in range(samples):
            t = i / self.sampling_rate
            # 基础正弦波
            value = amplitude * math.sin(2 * math.pi * base_freq * t)
            
            # 添加谐波增加真实感
            value += 0.3 * amplitude * math.sin(2 * math.pi * base_freq * 2 * t)
            value += 0.1 * amplitude * math.sin(2 * math.pi * base_freq * 3 * t)
            
            # 添加随机噪声
            noise = random.uniform(-0.1, 0.1) * amplitude
            value += noise
            
            signal.append(value)
        
        return signal
    
    def generate_multi_band_signal(self, duration: float) -> List[float]:
        """
        生成多频段混合的脑波信号
        
        Args:
            duration: 持续时间（秒）
            
        Returns:
            混合信号数据点列表
        """
        samples = int(duration * self.sampling_rate)
        mixed_signal = [0.0] * samples
        
        for band_name, (low_freq, high_freq) in self.frequency_bands.items():
            # 随机选择频段内的频率
            freq = random.uniform(low_freq, high_freq)
            # 随机振幅权重
            amplitude = random.uniform(0.5, 2.0)
            
            signal = self.generate_brain_wave_signal(duration, freq, amplitude)
            for i in range(samples):
                mixed_signal[i] += signal[i]
        
        # 归一化信号幅度
        max_abs = max(abs(x) for x in mixed_signal)
        if max_abs > 0:
            mixed_signal = [(x / max_abs) * settings.eeg_signal_max_amplitude for x in mixed_signal]
        
        return mixed_signal
    
    def generate_channel_data(self, duration: float) -> Dict[str, List[float]]:
        """
        生成所有通道的脑波数据
        
        Args:
            duration: 数据持续时间（秒）
            
        Returns:
            包含所有通道数据的字典
        """
        channel_data = {}
        
        for channel_name in self.channel_names:
            # 为每个通道生成略有不同的信号
            # 添加通道间的相位差和振幅差异
            phase_shift = random.uniform(0, 2 * math.pi)
            amplitude_factor = random.uniform(0.8, 1.2)
            
            # 生成基础多频段信号
            base_signal = self.generate_multi_band_signal(duration)
            
            # 应用通道特定的修改
            channel_signal = []
            for i, value in enumerate(base_signal):
                # 添加相位偏移
                t = i / self.sampling_rate
                phase_adjusted = value * math.cos(phase_shift)
                
                # 应用振幅因子
                final_value = phase_adjusted * amplitude_factor
                
                # 添加通道特定的小幅随机变化
                noise = random.uniform(-2, 2)
                final_value += noise
                
                channel_signal.append(round(final_value, 3))
            
            channel_data[channel_name] = channel_signal
        
        return channel_data
    
    def generate_30_second_data(self) -> List[Dict[str, Any]]:
        """
        生成30秒连续的脑波数据
        
        Returns:
            格式化的脑波数据列表，每个数据点包含时间戳和通道数据
        """
        duration = 30.0  # 30秒
        total_samples = int(self.sampling_rate * duration)
        
        # 生成所有通道的数据
        channel_data = self.generate_channel_data(duration)
        
        # 格式化为系统要求的数据格式
        formatted_data = []
        start_time = datetime.now(timezone.utc)
        
        for i in range(total_samples):
            # 计算当前时间戳
            current_time = start_time.timestamp() + (i / self.sampling_rate)
            
            # 构建数据点
            data_point = {
                "timestamp": current_time,
                "channels": {}
            }
            
            # 添加每个通道的数据
            for channel in self.channel_names:
                data_point["channels"][channel] = round(channel_data[channel][i], 3)
            
            formatted_data.append(data_point)
        
        return formatted_data
    
    def generate_batch_upload_data(self, id: str, room_id: str = None) -> Dict[str, Any]:
        """
        生成批量上传格式的数据
        
        Args:
            id: 会话ID
            room_id: 房间ID（可选）
            
        Returns:
            符合上传接口格式的数据
        """
        eeg_data = self.generate_30_second_data()
        
        upload_data = {
            "id": id,
            "data": eeg_data
        }
        
        if room_id:
            upload_data["room_id"] = room_id
        
        return upload_data
    
    def save_data_to_file(self, data: List[Dict[str, Any]], filename: str):
        """
        将数据保存到JSON文件
        
        Args:
            data: 要保存的数据
            filename: 文件名
        """
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
        
        print(f"数据已保存到: {filename}")
    
    def generate_analysis_result_data(self, id: str, room_id: str = None) -> Dict[str, Any]:
        """
        生成分析结果数据（模拟分析服务的输出）
        
        Args:
            id: 会话ID
            room_id: 房间ID（可选）
            
        Returns:
            分析结果数据
        """
        # 生成30秒的数据点
        data_points = []
        start_time = datetime.now(timezone.utc).timestamp()
        
        for i in range(30):  # 每秒一个数据点
            timestamp = start_time + i
            
            # 生成8频域数据
            frequency_data = {
                "delta_ratios": round(random.uniform(0.1, 0.3), 3),
                "theta_ratios": round(random.uniform(0.1, 0.25), 3),
                "low_alpha_ratios": round(random.uniform(0.05, 0.2), 3),
                "high_alpha_ratios": round(random.uniform(0.05, 0.2), 3),
                "low_beta_ratios": round(random.uniform(0.1, 0.25), 3),
                "high_beta_ratios": round(random.uniform(0.05, 0.2), 3),
                "low_gamma_ratios": round(random.uniform(0.02, 0.1), 3),
                "mid_gamma_ratios": round(random.uniform(0.02, 0.1), 3)
            }
            
            # 生成睡眠监测指标
            sleep_metrics = {
                "sleep_stage": random.choice([0, 1, 2, 3, 4]),  # 0-4睡眠阶段
                "sleep_quality": round(random.uniform(0.3, 0.9), 3),
                "attention_level": round(random.uniform(0.2, 0.8), 3),
                "relaxation_level": round(random.uniform(0.3, 0.9), 3)
            }
            
            data_point = {
                "timestamp": timestamp,
                **frequency_data,
                **sleep_metrics
            }
            
            data_points.append(data_point)
        
        result_data = {
            "id": id,
            "data": data_points
        }
        
        if room_id:
            result_data["room_id"] = room_id
        
        return result_data


if __name__ == "__main__":
    # 测试数据生成器
    generator = MockEEGDataGenerator()
    
    # 生成30秒数据
    print("正在生成30秒模拟脑波数据...")
    data = generator.generate_30_second_data()
    
    print(f"生成完成！共 {len(data)} 个数据点")
    print(f"采样率: {generator.sampling_rate} Hz")
    print(f"通道数: {generator.channels}")
    print(f"数据时长: 30 秒")
    
    # 保存示例数据
    generator.save_data_to_file(data[:10], "sample_eeg_data.json")  # 保存前10个数据点作为示例